// 测试 rivs:image_convert_to_topcc 与 OpenCV 实现一致性（从真实图片读取）
#include "image_converto.h"
#include <cstdio>
#include <cstdlib>
#include <opencv2/opencv.hpp>
#include <cmath>
#include <cassert>

extern "C" {
#include <libavutil/pixfmt.h>
}

bool compare_arrays(const float *arr1, const float *arr2, size_t n, float tol = 1e-4) {
    for (size_t i = 0; i < n; ++i) {
        if (std::fabs(arr1[i] - arr2[i]) > tol) return false;
    }
    return true;
}


int test_image_convert_to_L3()
{
    std::string image_path = "/root/test_videos/bus.jpg";
    cv::Mat src_img = cv::imread(image_path, cv::IMREAD_COLOR); // BGR
    if (src_img.empty()) {
        printf("Failed to read image: %s\n", image_path.c_str());
        return 3;
    }

    cv::Mat cv_img;
    float factor = MAX(src_img.cols/640.0f, src_img.rows/640.0f);
    cv::resize(src_img, cv_img, cv::Size(src_img.cols/factor, src_img.rows/factor));

    const int src_w = cv_img.cols;
    const int src_h = cv_img.rows;
    const int channel = 3;

    int dst_w = 640;
    int dst_h = 640;
    int device_id = 1;
    rivsSetDevice(device_id);
    // Device buffer for input AVFrame
    size_t image_size = static_cast<size_t>(src_h) * src_w * channel;
    uint8_t *image_dev = nullptr;
    if (rivsMalloc(&image_dev, image_size) != rivsSuccess) {
        printf("rivsMalloc failed for input image\n");
        return 4;
    }
    if (rivsMemcpy(image_dev, cv_img.data, image_size, rivsMemcpyHostToDevice) != rivsSuccess) {
        printf("rivsMemcpy H2D failed for input image\n");
        rivsFree(image_dev);
        return 5;
    }

    AVFrame image = AVFrame();
    image.data[0] = image_dev;
    image.format = AV_PIX_FMT_BGR24;
    image.height = src_h;
    image.width = src_w;

    // ROI 参数（以整张图为源，目标为 dst_w x dst_h）
    rivscv::ConvertParam param;
    param.m_srcROI.x = 0;
    param.m_srcROI.y = 0;
    param.m_srcROI.width = src_w;
    param.m_srcROI.height = src_h;
    param.m_dstROI.x = 0;
    param.m_dstROI.y = 0;
    param.m_dstROI.width = src_w;
    param.m_dstROI.height = src_h;
    param.m_swapRGB = true;
    param.m_scaleEnabled = true;
    param.m_scaleAlpha = 1.0f/255.0f;

    rivscv::RivsImageTensor out_tensor(device_id, {3, dst_h, dst_w}, RivsInference::DataType::TIF_FP32);
    int ret = rivscv::image_convert_to_L3(&image, &param, &out_tensor);
    if (ret != 0) {
        printf("image_convert_to_topcc failed: %d\n", ret);
        rivsFree(image_dev);
        return -1;
    }

    // Copy device result back
    const int out_width = dst_w;
    const int out_height = dst_h;
    size_t out_image_size = static_cast<size_t>(out_height) * out_width * channel * sizeof(float);
    std::vector<float> tops_result(out_height * out_width * channel);
    if (rivsMemcpy(tops_result.data(), out_tensor.m_data, out_image_size, rivsMemcpyDeviceToHost) != rivsSuccess) {
        printf("rivsMemcpy D2H failed for output tensor\n");
        rivsFree(image_dev);
        return -2;
    }

    // OpenCV reference pipeline
    cv::Rect roi(param.m_srcROI.x, param.m_srcROI.y, param.m_srcROI.width, param.m_srcROI.height);
    cv::Mat roi_img = cv_img(roi);
    cv::Mat out_img = cv::Mat::zeros(cv::Size(out_width, out_height), CV_8UC3);
    roi_img.copyTo(out_img(cv::Rect(0, 0, param.m_dstROI.width, param.m_dstROI.height)));
    cv::Mat rgb_img;
    cv::cvtColor(out_img, rgb_img, cv::COLOR_BGR2RGB);
    cv::Mat float_img;
    rgb_img.convertTo(float_img, CV_32FC3, param.m_scaleAlpha);
    // HWC->CHW
    std::vector<cv::Mat> chw(3);
    cv::split(float_img, chw);
    std::vector<float> opencv_result;
    for (int c = 0; c < 3; ++c) {
        opencv_result.insert(opencv_result.end(), (float*)chw[c].datastart, (float*)chw[c].dataend);
    }

    bool success = compare_arrays(tops_result.data(), opencv_result.data(), out_height * out_width * channel);
    printf("Compare with OpenCV: %s\n", success ? "SUCCESS" : "FAIL");

    // Optional: dump tops result back to a viewable JPEG for debugging
    // Convert CHW RGB float [0..1] to BGR8 HWC
    cv::Mat tops_bgr(out_height, out_width, CV_8UC3);
    const size_t hw = static_cast<size_t>(out_height) * out_width;
    for (int y = 0; y < out_height; ++y) {
        auto *row = tops_bgr.ptr<uint8_t>(y);
        for (int x = 0; x < out_width; ++x) {
            size_t i = static_cast<size_t>(y) * out_width + x;
            float r = tops_result[i];
            float g = tops_result[i + hw];
            float b = tops_result[i + 2 * hw];
            int ri = std::max(0, std::min(255, static_cast<int>(r * 255.0f + 0.5f)));
            int gi = std::max(0, std::min(255, static_cast<int>(g * 255.0f + 0.5f)));
            int bi = std::max(0, std::min(255, static_cast<int>(b * 255.0f + 0.5f)));
            row[x * 3 + 0] = static_cast<uint8_t>(bi);
            row[x * 3 + 1] = static_cast<uint8_t>(gi);
            row[x * 3 + 2] = static_cast<uint8_t>(ri);
        }
    }
    cv::imwrite("opencv_ref.jpg", out_img); // BGR
    cv::imwrite("rivs_result.jpg", tops_bgr);

    rivsFree(image_dev);
    return success ? 0 : 1;
}


int test_image_convert_to_L4()
{
    std::string image_path = "/root/test_videos/bus.jpg";
    cv::Mat src_img = cv::imread(image_path, cv::IMREAD_COLOR); // BGR
    if (src_img.empty()) {
        printf("Failed to read image: %s\n", image_path.c_str());
        return 3;
    }

    // Keep same preprocessing as L3 test
    cv::Mat cv_img;
    float factor = MAX(src_img.cols/640.0f, src_img.rows/640.0f);
    cv::resize(src_img, cv_img, cv::Size(src_img.cols/factor, src_img.rows/factor));

    const int src_w = cv_img.cols;
    const int src_h = cv_img.rows;
    const int channel = 3;

    int dst_w = 640;
    int dst_h = 640;
    int device_id = 1;

    // -- Prepare AVFrame for L3 path (device input) --
    size_t image_size = static_cast<size_t>(src_h) * src_w * channel;
    uint8_t *image_dev = nullptr;

    rivsSetDevice(device_id);

    if (rivsMalloc(&image_dev, image_size) != rivsSuccess) {
        printf("rivsMalloc failed for input image (L3)\n");
        return 4;
    }
    if (rivsMemcpy(image_dev, cv_img.data, image_size, rivsMemcpyHostToDevice) != rivsSuccess) {
        printf("rivsMemcpy H2D failed for input image (L3)\n");
        rivsFree(image_dev);
        return 5;
    }
    AVFrame image_dev_frame = AVFrame();
    image_dev_frame.data[0] = image_dev;
    image_dev_frame.format = AV_PIX_FMT_BGR24;
    image_dev_frame.height = src_h;
    image_dev_frame.width = src_w;

    // -- Prepare AVFrame for L4 path (host input) --
    AVFrame image_host_frame = AVFrame();
    image_host_frame.data[0] = cv_img.data;
    image_host_frame.format = AV_PIX_FMT_BGR24;
    image_host_frame.height = src_h;
    image_host_frame.width = src_w;

    // ROI and params
    rivscv::ConvertParam param;
    param.m_srcROI.x = 0; param.m_srcROI.y = 0; param.m_srcROI.width = src_w; param.m_srcROI.height = src_h;
    param.m_dstROI.x = 0; param.m_dstROI.y = 0; param.m_dstROI.width = src_w; param.m_dstROI.height = src_h;
    param.m_swapRGB = true;
    param.m_scaleEnabled = true;
    param.m_scaleAlpha = 1.0/255.0f;

    // Output tensors
    rivscv::RivsImageTensor out_l3(1, {3, dst_h, dst_w}, RivsInference::DataType::TIF_FP32);
    rivscv::RivsImageTensor out_l4(1,{3, dst_h, dst_w}, RivsInference::DataType::TIF_FP32);

    int ret = rivscv::image_convert_to_L3(&image_dev_frame, &param, &out_l3);
    if (ret != 0) {
        printf("image_convert_to_L3 failed: %d\n", ret);
        rivsFree(image_dev);
        return -1;
    }

    ret = rivscv::image_convert_to_L4(&image_host_frame, &param, &out_l4);
    if (ret != 0) {
        printf("image_convert_to_L4 failed: %d\n", ret);
        rivsFree(image_dev);
        return -2;
    }

    // Fetch both results to host and compare
    const int out_width = dst_w;
    const int out_height = dst_h;
    size_t out_image_size = static_cast<size_t>(out_height) * out_width * channel * sizeof(float);
    std::vector<float> l3(out_height * out_width * channel);
    std::vector<float> l4(out_height * out_width * channel);
    if (rivsMemcpy(l3.data(), out_l3.m_data, out_image_size, rivsMemcpyDeviceToHost) != rivsSuccess) {
        printf("D2H failed for L3\n");
        rivsFree(image_dev);
        return -3;
    }
    if (rivsMemcpy(l4.data(), out_l4.m_data, out_image_size, rivsMemcpyDeviceToHost) != rivsSuccess) {
        printf("D2H failed for L4\n");
        rivsFree(image_dev);
        return -4;
    }

    bool success = compare_arrays(l3.data(), l4.data(), out_height * out_width * channel, 1e-4f);
    printf("Compare L3 vs L4: %s\n", success ? "SUCCESS" : "FAIL");

    // Optional visualization (write L4 result)
    // Convert CHW float RGB [0..1] to BGR8 HWC for L4 result
    // cv::Mat l4_bgr(out_height, out_width, CV_8UC3);
    // const size_t hw = static_cast<size_t>(out_height) * out_width;
    // for (int y = 0; y < out_height; ++y) {
    //     auto *row = l4_bgr.ptr<uint8_t>(y);
    //     for (int x = 0; x < out_width; ++x) {
    //         size_t i = static_cast<size_t>(y) * out_width + x;
    //         float r = l4[i];
    //         float g = l4[i + hw];
    //         float b = l4[i + 2 * hw];
    //         int ri = std::max(0, std::min(255, static_cast<int>(r * 255.0f + 0.5f)));
    //         int gi = std::max(0, std::min(255, static_cast<int>(g * 255.0f + 0.5f)));
    //         int bi = std::max(0, std::min(255, static_cast<int>(b * 255.0f + 0.5f)));
    //         row[x * 3 + 0] = static_cast<uint8_t>(bi);
    //         row[x * 3 + 1] = static_cast<uint8_t>(gi);
    //         row[x * 3 + 2] = static_cast<uint8_t>(ri);
    //     }
    // }
    //cv::imwrite("l4_result.jpg", l4_bgr);

    rivscv::image_tensor_save(&out_l4, "l4_result.jpg");
    rivscv::image_tensor_save(&out_l3, "l3_result.jpg");

    // Convert CHW float RGB [0..1] to BGR8 HWC for L4 result
    // cv::Mat l3_bgr(out_height, out_width, CV_8UC3);
    // for (int y = 0; y < out_height; ++y) {
    //     auto *row = l4_bgr.ptr<uint8_t>(y);
    //     for (int x = 0; x < out_width; ++x) {
    //         size_t i = static_cast<size_t>(y) * out_width + x;
    //         float r = l3[i];
    //         float g = l3[i + hw];
    //         float b = l3[i + 2 * hw];
    //         int ri = std::max(0, std::min(255, static_cast<int>(r * 255.0f + 0.5f)));
    //         int gi = std::max(0, std::min(255, static_cast<int>(g * 255.0f + 0.5f)));
    //         int bi = std::max(0, std::min(255, static_cast<int>(b * 255.0f + 0.5f)));
    //         row[x * 3 + 0] = static_cast<uint8_t>(bi);
    //         row[x * 3 + 1] = static_cast<uint8_t>(gi);
    //         row[x * 3 + 2] = static_cast<uint8_t>(ri);
    //     }
    // }
    // cv::imwrite("l3_result.jpg", l3_bgr);

    rivsFree(image_dev);
    return success ? 0 : 1;
}


int main(int argc, char *argv[])
{
    test_image_convert_to_L3();
    test_image_convert_to_L4();

    return 0;
}
